import re

from pymongo import MongoClient
import pandas as pd


# 根据主题id获取中文名称
def get_classname(class_ids, mongo_ip, port, db, collection_classes):
    """
    根据主题编号获取主题中文名
    :param class_ids: 形如：["z0001","z0002","z0003"]
    :return: ["武器装备","战略思想","战史"]
    """

    class_list = []
    client = MongoClient(mongo_ip, port)
    collection = client[db][collection_classes]
    for class_id in class_ids:
        result = collection.find({"code": class_id})
        for i in result:
            name = i.get("name", " ")
            class_list.append(name)
    return class_list


def title_re(str_doc):
    """"0美军太空作战优势及弱点分析_李义_标签：作战理论_武器装备"""
    # 去掉字符串头部数字
    str_doc = re.sub(r"^\d+", "", str_doc)
    # 去除标签和作者
    str_doc = re.split(r'[: ：]', str_doc)[0].rsplit("_", 2)[0]
    return str_doc


# 正则对字符串清洗
def textParse(str_doc):
    # 正则过滤掉特殊符号、标点、英文、数字等。
    r1 = '[’!"#$%&\'*+,-./|<=>?@，—。?★、…【】《》？“”‘’· ！[\\]^_`{|}~]+'
    # 去除空格
    r2 = '\s+'
    str_doc = re.sub(r1, ' ', str_doc)
    # 多个空格成1个
    str_doc = re.sub(r2, ' ', str_doc)
    # 去除换行符
    str_doc = str_doc.replace('\n', ' ')
    str_doc = str_doc[0:100]
    return str_doc


def save_to_csv(mongo_ip, port, db, collection_articles,collection_classes, save_path):
    """
    将mongodb数据库中的表保存成csv文件，并对content字段截取30个字节
    :param MONGO_IP: mongodb数据库的ip
    :param PORT:  mongodb数据库的端口号
    :param DB: 要转换的数据库名
    :param COLLECTION: 要转换的数据库表名
    :param SAVE_PATH: csv文件的保存路径
    :return:
    """
    client = MongoClient(mongo_ip, port)
    collection = client[db][collection_articles]
    data = collection.find()
    data = list(data)
    df = pd.DataFrame(data)
    print(111, df['title'], type(df['title']))
    print(12, df['title'].str, type(df['title'].str))
    # df['title'] = title_re(df['title'].str) # 不能这样直接放到函数里，不是字符串类型
    # df['classes'] = get_classname(df['classes'])
    # df['content'] = textParse(df['content'].str)
    df['title'] = df['title'].apply(title_re)
    df['classes'] = df['classes'].apply(get_classname, args=(mongo_ip, port, db, collection_classes))
    df['content'] = df['content'].apply(textParse)
    # df['title']=df['title'].str.replace(r"^\d+","").rsplit("：|:", 1)[0].rsplit("_", 2)[0] #AttributeError: 'Series' object has no attribute 'rsplit'
    # df['title']=df['title'].str.replace(r"^\d+","").split("：|:")[0:-1].split("_")[0:-2]
    # df['classes']=df['classes']
    # df['content'] = df['content'].str.replace('[’!"#$%&\'*+,-./|<=>?@，—。?★、…【】《》？“”‘’· ！[\\]^_`{|}~]+',' ').replace('\s+',' ').replace('\n', ' ')[0:100]
    df.to_csv(SAVE_PATH, index=False, mode='w', encoding="utf_8_sig")  # 存储格式为 utf-8


if __name__ == '__main__':
    MONGO_IP = '127.0.0.1'
    PORT = 27017
    DB = "lc_search"
    COLLECTION_ARTICLES= "articles"
    COLLECTION_CLASSES= "classes"
    SAVE_PATH = "D:\data_csv\data204.csv"
    save_to_csv(MONGO_IP, PORT, DB, COLLECTION_ARTICLES,COLLECTION_CLASSES, SAVE_PATH)
